feat(sc): rollout path — TQReplayBuffer, rollout_manager, _rollout_pump#3219
Conversation
yuki-97
added
the
CI:Lfast
|
/ok to test a14ddff |
|
/ok to test c6ad118 |
02ca7f0 to
97ce826
Compare
|
/ok to test 97ce826 |
|
/ok to test bf73365 |
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Team review of the rollout-path infrastructure (PR 1/4 of the #2819 split). The split is clean: parity coverage from test_rollouts.py moved 1:1 into test_rollout_manager.py (both *_matches_original guards survive), removed symbols (ReplayBufferNew, _rollout_done, _flush_incomplete_groups) have no dangling references, and the producer→consumer DataPlane contract (sample-id scheme + weight_version tag) stays compatible with the untouched _train_pump. Linters pass; 18 of the PR's unit tests were run locally on GPU (all pass), including test_async_rollout_manager_matches_original.
Findings were cross-checked against the rest of the stack (#3220 / #3266 / #3267) — items already resolved downstream (e.g. the dead max_rollout_prompts field, deleted in #3220; the finite-max_num_epochs termination gap, closed by #3220's dead-code removal plus #3266's max_num_steps ≤ max_num_epochs × len(dataloader) clamp) are intentionally NOT flagged inline. The remaining inline comments cover only what survives the full stack: a silent-failure/permit-leak path in the new fire-and-forget rollout dispatch, a missing @pytest.mark.vllm, stale excluded_unit_tests.sh deselect paths after the test move, an untested default dispatch branch, a likely-rebase-artifact pyrefly.toml removal, and a doc nit.
Non-blocking observations (no action needed in this PR):
RolloutManager'smax_rollout_turnsdefault changed fromNone(coerced to 999999 on the native path) to1, and the coercion was removed. No production caller exists until PR#3, but a future native multi-turn caller that omits the arg now silently gets single-turn rollouts — worth keeping in mind for the PR#3 wiring.- The new config-validation raises in
SingleControllerActor.__init__live under# pragma: no cover; factoring them into a module-level helper would make them unit-testable (optional).
Generated by Claude Code
bf73365 to
e48e80b
Compare
|
/ok to test efc7d31 |
efc7d31 to
c454acf
Compare
|
/ok to test c454acf |
c454acf to
f929aa0
Compare
|
/ok to test f929aa0 |
Signed-off-by: Yuki Huang <yukih@nvidia.com>
…add target_step to _FakeBuffer.reserve Signed-off-by: Yuki Huang <yukih@nvidia.com>
…shutdown path Signed-off-by: Yuki Huang <yukih@nvidia.com>
…manager signatures Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
…force_in_order Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: ruit <ruit@nvidia.com>
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Follow-up review of the commits addressing the earlier round (bf73365 → f929aa0). The fixes hold up well under an adversarial concurrency pass: permit accounting is exactly-once on every interleaving we could construct (including cancel-before-start via task_started_event), run() now propagates rollout failures without orphaning tasks, and commit()'s rollback has no window against the train pump's claim path. The new tests genuinely exercise these paths, and the added @pytest.mark.vllm also fixes a latent lane bug (the test previously ran GPU work in the CPU-only Other shard). One low-severity note inline.
Generated by Claude Code
Signed-off-by: ruit <ruit@nvidia.com>
f929aa0 to
48184ad
Compare
|
/ok to test 48184ad |
Signed-off-by: ruit <ruit@nvidia.com>
|
/ok to test 62a759a |
3 participants
Part 1/4 of the #2819 split. See #2819 (reference) or #3267 (4/4) for the full code and context.
#3219 (1/4) -> #3220 (2/4) -> #3266 (3/4) -> #3267 (4/4)
Summary
Rollout-side infrastructure for the SC (Single Controller) async-GRPO training path.
TQReplayBuffer(nemo_rl/algorithms/async_utils/replay_buffer.py): group-granular replay buffer with reserve/commit slot accounting; producer-side tensorization writes N training-shaped rows per prompt group. Replaces the WIPReplayBufferNew.RolloutManager.generate_and_push(nemo_rl/experience/rollout_manager.py): reserve a buffer slot, run one prompt rollout, commit with start/end weight versions. Also renamestask_to_env → env_handlesacross bothAsyncRolloutImplandAsyncNemoGymRolloutImpl.payload.py(nemo_rl/experience/payload.py): newpack_payload+record_to_train_batchhelpers used byTQReplayBuffer.commit.SingleControllerActor._rollout_pump: rewritten to dispatch prompts through the newRolloutManager, with per-weight-version quota (over_sampling=false) and exact target-step matching (force_in_order=true). NewSingleControllerConfigfieldsover_sampling/force_in_order; actor__init__gainsrollout_manager/dataloaderfor driver-side wiring.test_tq_replay_bufferandtest_rollout_pump; newtest_rollout_managerreplaces the oldtest_rollouts.py;TestReplayBufferNewremoved fromtest_async_utils.Note:
_train_pumpis left in its pre-refactor form in this PR (still usesdp_client.claim_meta+ the old top-levelStalenessSampler); PR#2 switches it to the newsampler.evict/selectAPI. SC becomes end-to-end runnable only after PR#3 lands the driver-side setup and entrypoint.Usage
min_groups_for_streaming_train(current calledmin_groups_per_batch, will be renamed in later PR)Known Missing Features (Tracked in #2625)
max_inflight_prompts->max_inflight_rolloutsover_sampling_ratio(>1) to avoid over sample to much and waste data.